package com.example.base;

import com.example.entity.Customer;
import org.apache.commons.lang3.StringUtils;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.ReduceFunction;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
import org.apache.flink.connector.datagen.source.DataGeneratorSource;
import org.apache.flink.connector.datagen.source.GeneratorFunction;
import org.apache.flink.streaming.api.datastream.*;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.windowing.AllWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.TumblingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;

import java.time.Duration;
import java.time.LocalDateTime;
import java.time.ZoneOffset;
import java.util.*;

/**
 * Author wangJinLong
 * Date 2025/8/12 11:11
 **/
public class FlinkStreamDataGen {
    private static final String FORMAT = "yyyy-MM-dd HH:mm:ss";
    private static List<String> faceList = Arrays.asList("14901", "24902", "34903", "44903", "54903");
    private static List<String> siteList = Arrays.asList("05350101", "05350104", "05350122", "00004", "00002");

    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(3);
        Random random = new Random();

        GeneratorFunction<Long, Customer> generatorFunction = value -> {
            Customer customer = new Customer();
            customer.setSiteSn(siteList.get(random.nextInt(5)));
            customer.setFeaturesId(faceList.get(random.nextInt(5)));
            LocalDateTime now = LocalDateTime.now();
            long epochMilli = now.toInstant(ZoneOffset.systemDefault().getRules().getOffset(now)).toEpochMilli();
            customer.setCaptureTime(epochMilli);
            return customer;
        };

        DataGeneratorSource<Customer> source = new DataGeneratorSource<>(
                generatorFunction,
                60L,
                RateLimiterStrategy.perSecond(1),
                Types.POJO(Customer.class)
        );

        DataStream<Customer> dataStream = env.fromSource(
                source,
                WatermarkStrategy.
                        <Customer>forBoundedOutOfOrderness(Duration.ofSeconds(5))
                        .withTimestampAssigner((event, timestamp) -> event.getCaptureTime()),
                "source"
        );

        KeyedStream<Customer, String> keyedStream = dataStream.keyBy(Customer::getFeaturesId);

        DataStream<Customer> reducedStream = keyedStream.reduce(new ReduceFunction<Customer>() {
            @Override
            public Customer reduce(Customer value1, Customer value2) throws Exception {
                Customer customer = value1.getCaptureTime() > value2.getCaptureTime() ? value1 : value2;
                return customer;
            }
        });

        reducedStream.print();

        env.execute();
    }

    private static void extracted4() throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        Random random = new Random();

        GeneratorFunction<Long, Customer> generatorFunction = value -> {
            Customer customer = new Customer();
            customer.setSiteSn(siteList.get(random.nextInt(5)));
            customer.setFeaturesId(faceList.get(random.nextInt(5)));
            LocalDateTime now = LocalDateTime.now();
            long epochMilli = now.toInstant(ZoneOffset.systemDefault().getRules().getOffset(now)).toEpochMilli();
            customer.setCaptureTime(epochMilli);
            return customer;
        };

        DataGeneratorSource<Customer> source = new DataGeneratorSource<>(
                generatorFunction,
                100L,
                RateLimiterStrategy.perSecond(1),
                Types.POJO(Customer.class)
        );

        DataStream<Customer> dataStream = env.fromSource(
                source,
                WatermarkStrategy.
                        <Customer>forBoundedOutOfOrderness(Duration.ofSeconds(5))
                        .withTimestampAssigner((event, timestamp) -> event.getCaptureTime()),
                "source"
        );

        AllWindowedStream<Customer, TimeWindow> windowedStream = dataStream
                .windowAll(TumblingEventTimeWindows.of(Duration.ofSeconds(10)));

        DataStream<String> windowStream = windowedStream.apply(new AllWindowFunction<>() {
            @Override
            public void apply(TimeWindow window, Iterable<Customer> values, Collector<String> out) {
                Set<String> siteSns = new HashSet<>();
                for (Customer value : values) {
                    siteSns.add(value.getSiteSn());
                }
                String join = StringUtils.joinWith(", ", siteSns);
                out.collect(join); // 无法使用 list 或 set 集合
            }
        });

        windowStream.print();

        env.execute();
    }

    private static void extracted3() throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        GeneratorFunction<Long, Customer> generatorFunction = value -> {
            Customer customer = new Customer();
            customer.setSiteSn("siteSn: " + value);
            customer.setFeaturesId("featuresId: " + value);
            LocalDateTime now = LocalDateTime.now();
            long epochMilli = now.toInstant(ZoneOffset.systemDefault().getRules().getOffset(now)).toEpochMilli();
            customer.setCaptureTime(epochMilli);
            return customer;
        };

        DataGeneratorSource<Customer> source = new DataGeneratorSource<>(
                generatorFunction,
                100L,
                RateLimiterStrategy.perSecond(1),
                Types.POJO(Customer.class)
        );

        DataStream<Customer> dataStream = env.fromSource(
                source,
                WatermarkStrategy.
                        <Customer>forBoundedOutOfOrderness(Duration.ofSeconds(5))
                        .withTimestampAssigner((event, timestamp) -> event.getCaptureTime()),
                "source"
        );

        AllWindowedStream<Customer, TimeWindow> windowedStream = dataStream
                .windowAll(TumblingEventTimeWindows.of(Duration.ofSeconds(10)));

        SingleOutputStreamOperator<String> windowStream = windowedStream.apply(new AllWindowFunction<>() {
            @Override
            public void apply(TimeWindow window, Iterable<Customer> values, Collector<String> out) {
                StringBuilder name = new StringBuilder("\n");
                values.forEach(value -> name.append(value.getSiteSn()).append(", "));
                out.collect(name.toString());
            }
        });

        windowStream.print();

        env.execute();
    }

    /**
     * 控制数据生成速度
     */
    private static void extracted2() throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        GeneratorFunction<Long, Customer> generatorFunction = value -> {
            Customer customer = new Customer();
            customer.setSiteSn("siteSn: " + value);
            customer.setFeaturesId("featuresId: " + value);
            LocalDateTime now = LocalDateTime.now();
            long captureTime = now.toInstant(ZoneOffset.systemDefault().getRules().getOffset(now)).getEpochSecond();
            customer.setCaptureTime(captureTime);
            return customer;
        };

        DataGeneratorSource<Customer> source = new DataGeneratorSource<>(
                generatorFunction,
                100L,
                RateLimiterStrategy.perSecond(1),
                Types.POJO(Customer.class)
        );

        DataStream<Customer> dataStream = env.fromSource(source, WatermarkStrategy.noWatermarks(), "source");

        dataStream.print();

        env.execute();
    }

    /**
     * 生成固定数量之后结束
     * @throws Exception
     */
    private static void extracted1() throws Exception {
        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();
        environment.setParallelism(1);

        GeneratorFunction<Long, String> generatorFunction = index -> "Number: " + index;
        long numberOfRecords = 100;

        DataGeneratorSource<String> source =
                new DataGeneratorSource<>(generatorFunction, numberOfRecords, Types.STRING);

        DataStream<String> stream =
                environment.fromSource(
                        source,
                        WatermarkStrategy.noWatermarks(),
                        "Generator Source"
                );

        stream.print();
        environment.execute();
    }

    // 自定义生成的数据类型
    private static void extracted() throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);

        GeneratorFunction<Long, Customer> generatorFunction1 = value -> {
            Customer customer = new Customer();
            customer.setSiteSn("siteSn: " + value);
            customer.setFeaturesId("featuresId: " + value);
            LocalDateTime now = LocalDateTime.now();
            long captureTime = now.toInstant(ZoneOffset.systemDefault().getRules().getOffset(now)).getEpochSecond();
            customer.setCaptureTime(captureTime);
            return customer;
        };
        DataGeneratorSource<Customer> source =
                new DataGeneratorSource<>(generatorFunction1, 100L, Types.POJO(Customer.class));

        DataStream<Customer> swsss = env.fromSource(source, WatermarkStrategy.noWatermarks(), "swsss");

        swsss.print();

        env.execute();
    }
}
